"""
Milvus客户端 - 向量存储和检索
"""
import logging
import time
import json
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from pymilvus import (
    connections, Collection, CollectionSchema, FieldSchema, DataType,
    utility, MilvusException
)

from ..config.settings import MILVUS_CONFIG, VECTOR_STORE_CONFIG

logger = logging.getLogger(__name__)

class MilvusClient:
    """Milvus客户端"""
    
    def __init__(self):
        self.host = MILVUS_CONFIG["host"]
        self.port = int(MILVUS_CONFIG["port"])  # 确保端口是整数
        self.user = MILVUS_CONFIG["user"]
        self.password = MILVUS_CONFIG["password"]
        self.db_name = MILVUS_CONFIG["db_name"]
        
        self.collection_name = VECTOR_STORE_CONFIG["collection_name"]
        self.dimension = VECTOR_STORE_CONFIG["dimension"]  # 使用配置中的维度
        self.index_type = VECTOR_STORE_CONFIG["index_type"]
        self.metric_type = VECTOR_STORE_CONFIG["metric_type"]
        
        self.connection = None
        self.collection = None
        
        # 连接Milvus
        self._connect()
    
    def _connect(self):
        """连接到Milvus"""
        try:
            # 连接参数 - 确保端口是字符串格式
            connection_params = {
                "alias": "default",
                "host": self.host,
                "port": str(self.port)  # pymilus期望端口是字符串
            }
            
            if self.user and self.password:
                connection_params.update({
                    "user": self.user,
                    "password": self.password
                })
            
            # 添加超时设置
            connection_params["timeout"] = 10
            
            logger.info(f"尝试连接到Milvus: {self.host}:{self.port}")
            logger.info(f"连接参数: {connection_params}")
            
            # 建立连接
            connections.connect(**connection_params)
            
            # 验证连接
            if connections.has_connection("default"):
                # 保存连接对象
                self.connection = connections._fetch_handler("default")
                logger.info(f"成功连接到Milvus: {self.host}:{self.port}")
                
                # 使用指定数据库
                if self.db_name and self.db_name != "default":
                    utility.using_database(self.db_name)
                    logger.info(f"使用数据库: {self.db_name}")
            else:
                logger.error("Milvus连接验证失败")
                self.connection = None
                
        except Exception as e:
            logger.error(f"连接Milvus失败: {e}")
            self.connection = None  # 确保连接为None而不是抛出异常
    
    def _create_collection_schema(self) -> CollectionSchema:
        """创建集合模式"""
        fields = [
            # 主键字段
            FieldSchema(
                name="id",
                dtype=DataType.INT64,
                is_primary=True,
                auto_id=True
            ),
            # 向量字段
            FieldSchema(
                name="embedding",
                dtype=DataType.FLOAT_VECTOR,
                dim=self.dimension
            ),
            # 文本内容
            FieldSchema(
                name="content",
                dtype=DataType.VARCHAR,
                max_length=65535
            ),
            # 文件ID
            FieldSchema(
                name="file_id",
                dtype=DataType.VARCHAR,
                max_length=64
            ),
            # 块ID
            FieldSchema(
                name="chunk_id",
                dtype=DataType.VARCHAR,
                max_length=64
            ),
            # 块索引
            FieldSchema(
                name="chunk_index",
                dtype=DataType.INT64
            ),
            # 文件类型
            FieldSchema(
                name="file_type",
                dtype=DataType.VARCHAR,
                max_length=20
            ),
            # 文件名
            FieldSchema(
                name="filename",
                dtype=DataType.VARCHAR,
                max_length=255
            ),
            # 元数据（JSON格式）
            FieldSchema(
                name="metadata",
                dtype=DataType.VARCHAR,
                max_length=65535
            ),
            # 创建时间
            FieldSchema(
                name="created_time",
                dtype=DataType.INT64
            )
        ]
        
        schema = CollectionSchema(
            fields=fields,
            description="文档向量存储集合"
        )
        
        return schema
    
    def create_collection(self, collection_name: str = None) -> bool:
        """
        创建集合
        
        Args:
            collection_name: 集合名称
            
        Returns:
            是否创建成功
        """
        try:
            collection_name = collection_name or self.collection_name
            
            # 检查集合是否已存在
            if utility.has_collection(collection_name):
                # 检查现有集合的维度是否与配置一致
                collection = Collection(collection_name)
                schema = collection.schema
                embedding_field = None
                
                # 查找embedding字段
                for field in schema.fields:
                    if field.name == "embedding":
                        embedding_field = field
                        break
                
                # 如果维度不匹配，需要重新创建集合
                if embedding_field and embedding_field.params.get("dim") != self.dimension:
                    logger.warning(f"集合维度不匹配: 期望 {self.dimension}, 实际 {embedding_field.params.get('dim')}")
                    logger.info("删除现有集合以重新创建")
                    collection.drop()
                else:
                    logger.info(f"集合已存在: {collection_name}")
                    return True
            
            # 创建集合模式
            schema = self._create_collection_schema()
            
            # 创建集合
            collection = Collection(
                name=collection_name,
                schema=schema,
                using="default"  # 使用默认连接别名
            )
            
            logger.info(f"集合创建成功: {collection_name}，维度: {self.dimension}")
            
            # 创建索引
            self._create_index(collection)
            
            return True
            
        except Exception as e:
            logger.error(f"创建集合失败: {e}")
            logger.exception(e)
            return False
    
    def _create_index(self, collection: Collection):
        """创建索引"""
        try:
            # 索引参数
            if self.index_type == "HNSW":
                index_params = {
                    "metric_type": self.metric_type,
                    "index_type": self.index_type,
                    "params": {
                        "M": VECTOR_STORE_CONFIG["m"],
                        "efConstruction": VECTOR_STORE_CONFIG["efConstruction"]
                    }
                }
            elif self.index_type == "IVF_FLAT":
                index_params = {
                    "metric_type": self.metric_type,
                    "index_type": self.index_type,
                    "params": {
                        "nlist": VECTOR_STORE_CONFIG["nlist"]
                    }
                }
            else:
                # 默认使用FLAT索引
                index_params = {
                    "metric_type": self.metric_type,
                    "index_type": "FLAT"
                }
            
            # 创建索引
            collection.create_index(
                field_name="embedding",
                index_params=index_params
            )
            
            logger.info(f"索引创建成功: {self.index_type}")
            
        except Exception as e:
            logger.error(f"创建索引失败: {e}")
            raise
    
    def get_collection(self, collection_name: str = None) -> Optional[Collection]:
        """
        获取集合
        
        Args:
            collection_name: 集合名称
            
        Returns:
            集合对象
        """
        try:
            collection_name = collection_name or self.collection_name
            
            if not utility.has_collection(collection_name):
                logger.warning(f"集合不存在: {collection_name}")
                return None
            
            collection = Collection(
                name=collection_name,
                using="default"  # 使用默认连接别名
            )
            
            return collection
            
        except Exception as e:
            logger.error(f"获取集合失败: {e}")
            return None
    
    def insert_vectors(self, vectors_data: List[Dict[str, Any]], collection_name: str = None) -> List[int]:
        """
        插入向量数据
        
        Args:
            vectors_data: 向量数据列表
            collection_name: 集合名称
            
        Returns:
            插入的ID列表
        """
        try:
            logger.info(f"开始插入 {len(vectors_data)} 条向量数据到Milvus")
            
            collection = self.get_collection(collection_name)
            if not collection:
                error_msg = f"无法获取集合: {collection_name or self.collection_name}"
                logger.error(error_msg)
                return []
            
            # 加载集合
            logger.debug("加载集合")
            collection.load()
            
            # 准备数据
            logger.debug("准备插入数据")
            embeddings = []
            contents = []
            file_ids = []
            chunk_ids = []
            chunk_indices = []
            file_types = []
            filenames = []
            metadatas = []
            created_times = []
            
            current_time = int(time.time())
            
            for i, vector_data in enumerate(vectors_data):
                try:
                    embeddings.append(vector_data["embedding"])
                    contents.append(vector_data["content"])
                    file_ids.append(vector_data["file_id"])
                    chunk_ids.append(vector_data["chunk_id"])
                    chunk_indices.append(vector_data["chunk_index"])
                    file_types.append(vector_data["file_type"])
                    filenames.append(vector_data["filename"])
                    
                    # 处理元数据
                    metadata = vector_data.get("metadata", {})
                    if isinstance(metadata, dict):
                        metadatas.append(json.dumps(metadata, ensure_ascii=False))
                    else:
                        metadatas.append(str(metadata))
                    
                    created_times.append(current_time)
                except KeyError as e:
                    logger.error(f"向量数据 {i} 缺少必要字段: {e}")
                    logger.debug(f"向量数据内容: {vector_data}")
                    continue
            
            # 插入数据
            insert_data = [
                embeddings,
                contents,
                file_ids,
                chunk_ids,
                chunk_indices,
                file_types,
                filenames,
                metadatas,
                created_times
            ]
            
            logger.debug(f"执行插入操作，数据维度: {[len(d) for d in insert_data]}")
            result = collection.insert(insert_data)
            
            # 刷新集合
            logger.debug("刷新集合")
            collection.flush()
            
            logger.info(f"成功插入 {len(vectors_data)} 条向量数据，插入ID数量: {len(result.primary_keys)}")
            return result.primary_keys
            
        except Exception as e:
            error_msg = f"插入向量数据失败: {e}"
            logger.error(error_msg)
            logger.exception(e)  # 打印完整的异常堆栈
            return []
    
    def search_vectors(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        collection_name: str = None,
        filter_expr: str = None
    ) -> List[Dict[str, Any]]:
        """
        搜索向量
        
        Args:
            query_embedding: 查询向量
            top_k: 返回结果数量
            collection_name: 集合名称
            filter_expr: 过滤表达式
            
        Returns:
            搜索结果列表
        """
        try:
            logger.info(f"开始执行向量搜索: top_k={top_k}, collection_name={collection_name}")
            logger.debug(f"查询向量维度: {len(query_embedding)}")
            
            collection = self.get_collection(collection_name)
            if not collection:
                logger.warning("无法获取集合，返回空结果")
                return []
            
            # 检查查询向量维度
            if len(query_embedding) != self.dimension:
                logger.error(f"查询向量维度不匹配: 期望 {self.dimension}, 实际 {len(query_embedding)}")
                return []
            
            # 加载集合
            logger.debug("加载集合")
            collection.load()
            
            # 搜索参数
            search_params = {
                "metric_type": self.metric_type,
                "params": {}
            }
            
            if self.index_type == "HNSW":
                search_params["params"]["ef"] = VECTOR_STORE_CONFIG["ef"]
            elif self.index_type == "IVF_FLAT":
                search_params["params"]["nprobe"] = 10
            
            logger.info(f"执行向量搜索: top_k={top_k}, metric_type={self.metric_type}, dimension={len(query_embedding)}")
            logger.debug(f"搜索参数: {search_params}")
            logger.debug(f"过滤表达式: {filter_expr}")
            
            # 执行搜索
            logger.debug("开始执行Milvus搜索")
            results = collection.search(
                data=[query_embedding],
                anns_field="embedding",
                param=search_params,
                limit=top_k,
                expr=filter_expr,
                output_fields=["content", "file_id", "chunk_id", "chunk_index", 
                             "file_type", "filename", "metadata", "created_time"]
            )
            
            logger.debug(f"搜索完成，原始结果数量: {len(results)}")
            
            # 处理结果
            search_results = []
            for i, hits in enumerate(results):
                logger.debug(f"处理第 {i+1} 组结果，包含 {len(hits)} 个匹配项")
                for hit in hits:
                    # 解析元数据
                    metadata_str = hit.entity.metadata if hasattr(hit.entity, 'metadata') else "{}"
                    try:
                        metadata = json.loads(metadata_str) if isinstance(metadata_str, str) else metadata_str
                    except:
                        metadata = {}
                    
                    result = {
                        "id": hit.id,
                        "score": hit.score,
                        "content": hit.entity.content if hasattr(hit.entity, 'content') else "",
                        "file_id": hit.entity.file_id if hasattr(hit.entity, 'file_id') else "",
                        "chunk_id": hit.entity.chunk_id if hasattr(hit.entity, 'chunk_id') else "",
                        "chunk_index": hit.entity.chunk_index if hasattr(hit.entity, 'chunk_index') else 0,
                        "file_type": hit.entity.file_type if hasattr(hit.entity, 'file_type') else "",
                        "filename": hit.entity.filename if hasattr(hit.entity, 'filename') else "",
                        "metadata": metadata,
                        "created_time": hit.entity.created_time if hasattr(hit.entity, 'created_time') else 0
                    }
                    search_results.append(result)
                    logger.debug(f"结果: ID={hit.id}, Score={hit.score:.4f}, 内容预览: {result['content'][:50]}")
            
            logger.info(f"总共返回 {len(search_results)} 个结果")
            return search_results
            
        except Exception as e:
            logger.error(f"搜索向量失败: {e}")
            logger.exception(e)
            return []
    
    def delete_vectors(self, file_id: str, collection_name: str = None) -> bool:
        """
        删除指定文件的向量
        
        Args:
            file_id: 文件ID
            collection_name: 集合名称
            
        Returns:
            是否删除成功
        """
        try:
            collection = self.get_collection(collection_name)
            if not collection:
                return False
            
            # 删除条件
            expr = f'file_id == "{file_id}"'
            
            # 执行删除
            collection.delete(expr)
            
            # 刷新集合
            collection.flush()
            
            logger.info(f"成功删除文件 {file_id} 的向量数据")
            return True
            
        except Exception as e:
            logger.error(f"删除向量失败: {e}")
            return False
    
    def get_collection_stats(self, collection_name: str = None) -> Dict[str, Any]:
        """
        获取集合统计信息
        
        Args:
            collection_name: 集合名称
            
        Returns:
            统计信息
        """
        try:
            collection = self.get_collection(collection_name)
            if not collection:
                return {}
            
            # 获取集合信息
            stats = {
                "name": collection.name,
                "num_entities": collection.num_entities,
                "schema": collection.schema.to_dict()
            }
            
            # 获取索引信息
            index_info = collection.index().params
            stats["index"] = index_info
            
            return stats
            
        except Exception as e:
            logger.error(f"获取集合统计信息失败: {e}")
            return {}
    
    def drop_collection(self, collection_name: str = None) -> bool:
        """
        删除集合
        
        Args:
            collection_name: 集合名称
            
        Returns:
            是否删除成功
        """
        try:
            collection_name = collection_name or self.collection_name
            
            if not utility.has_collection(collection_name):
                logger.warning(f"集合不存在: {collection_name}")
                return True
            
            utility.drop_collection(collection_name)
            logger.info(f"成功删除集合: {collection_name}")
            return True
            
        except Exception as e:
            logger.error(f"删除集合失败: {e}")
            return False
    
    def close(self):
        """关闭连接"""
        try:
            if self.connection:
                connections.disconnect(self.connection.alias)
                logger.info("Milvus连接已关闭")
        except Exception as e:
            logger.error(f"关闭连接失败: {e}")

# 全局实例
milvus_client = MilvusClient()
